package com.szsh.bigdata.open.utils

import org.apache.spark.sql.Row
import org.apache.spark.sql.execution.streaming.StreamMetadata.format
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, StringType, StructField, StructType}
import org.json4s.jackson.Serialization.write

/**
 * 用户自定义Agg函数，用于合并所有字符串
 */
object StringMergeDistinct extends UserDefinedAggregateFunction {
	override def inputSchema: StructType = StructType(Array(StructField("input", StringType)))
	
	override def bufferSchema: StructType = StructType(Array(StructField("count", StringType)))
	
	override def dataType: DataType = StringType
	
	override def deterministic: Boolean = true
	
	override def initialize(buffer: MutableAggregationBuffer): Unit = {
		buffer(0) = null
	}
	
	/**
	 * 更新中间结果, 将每次传过来的dataframe的一行添加到分片的中间结果中
	 *
	 * @param buffer 整个分片遍历过来的中间结果
	 * @param input  dataframe的一行
	 */
	override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
		//    buffer.getMap(0)
		buffer(0) = buffer.getString(0) + ",," + input
	}
	
	/**
	 * 分片的合并, 将每次传过来的分片的添加到总分片的中间结果中
	 *
	 * @param buffer1 整个合并过程的中间结果
	 * @param buffer2 一个分片的中间结果
	 */
	override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
		buffer1(0) = buffer1.getString(0) + ",," + buffer2.getString(0)
	}
	
	/**
	 * 返回函数结果,用";;"分割不同字符串
	 *
	 * @param buffer 并过程的中间结果buffer遍历所有分片结束后的结果
	 * @return
	 */
	override def evaluate(buffer: Row): Any =
		write(buffer.getString(0)
			.replaceAll("\\[", "").replaceAll("\\]", "")
			.split(",,").filter(_ != "null").distinct.toList.sortBy(x => x))
}
